// Copyright (C) 2011  Andrew H. Chan
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

#ifndef _APPROX_H__
#define _APPROX_H__

#include "typedefs.h"
#include "graph.h"

template<class T, class S>
struct pair_first_less {
    bool operator()(std::pair<T, S> const& a, std::pair<T, S> const&b) const {
        if(a.first == b.first) {
            return a.second < b.second;
        }
        return a.first < b.first;
    }
};

template<class T, class S, class R>
struct triple_first_less {
    bool operator()(boost::tuple<T, S, R> const& a, boost::tuple<T, S, R> const& b) const {
        if(a.template get<0>() == b.template get<0>()) {
            if(a.template get<1>() == b.template get<1>()) {
                return a.template get<2>() < b.template get<2>();
            }
            return a.template get<1>() < b.template get<1>();
        }
        return a.template get<0>() < b.template get<0>();
        return false;
    }
};

class SteinerApprox {
    public:
        template<class T> Graph<T> steiner_approx(Graph<T> const& graph, T root, boost::unordered_set<T> terminal_set, int i, int k);
        template<class T> std::vector<std::pair<T, char> > shortest_path(Graph<T> const& graph, T source, T dest);
        // unordered_map<T, T>:  vertex -> previous vertex
        // unordered_map<T, double>: vertex -> distance from source 
        template<class T> int dijkstra(Graph<T> const& graph, T source, boost::unordered_map<T,double>& distances, boost::unordered_map<T,std::pair<T,char> >& prev);
        template<class T> int merge_trees(Graph<T>& graph1, Graph<T> const& graph2);
};

template<class T> Graph<T> SteinerApprox::steiner_approx(Graph<T> const& graph, T root, boost::unordered_set<T> terminal_set, int i, int k_orig) {
    //assert(k == (int)terminal_set.size());
    // remove root from terminal set
    //terminal_set.erase(root);

    int k = k_orig;
    // there do not exist k terminals to connect from root
    if((int)terminal_set.size() < k) {
        return Graph<T>(); // return empty graph
    }
    // if the root is a dummy terminal, then return empty graph
    if(graph.terminal_set.find(root) != graph.terminal_set.end()){
        return Graph<T>(); // return the empty graph
    }

    if(i == 1) {
        Graph<T> return_graph;
        boost::unordered_map<T, double> distances;
        boost::unordered_map<T, std::pair<T, char> > prev;
        // run Dijkstra from root
        dijkstra(graph, root, distances, prev);
        std::vector<std::pair<double, T> > list_of_terminals;
        // iterate over terminals and take the k closest
        for(typename boost::unordered_set<T>::const_iterator term_iter = terminal_set.begin(); term_iter != terminal_set.end(); ++term_iter) {
            // find k closest, reachable terminals
            if(distances.find(*term_iter) != distances.end()) {
                list_of_terminals.push_back(std::pair<double, T>(distances[*term_iter], *term_iter));
            }
        }
        
        // there do not exist k terminals reachable from the root
        if((int)list_of_terminals.size() < k) {
            return return_graph; // return the empty graph
        }
        //assert((int)list_of_terminals.size() >= k);
        std::sort(list_of_terminals.begin(), list_of_terminals.end(), pair_first_less<double,T>());
        std::vector<T> closest_terminals; 
        for(int ii=0; ii < k; ++ii) {
            closest_terminals.push_back(list_of_terminals[ii].second);
        }
        // add paths from root to each terminal
        for(uint ii=0; ii<closest_terminals.size(); ++ii) {
            // generate path from root to terminal ii
            std::vector<std::pair<T, char> > path;
            T tmp_vert = closest_terminals[ii];
            path.push_back(std::pair<T,char>(tmp_vert, -1)); // push dest
            while(tmp_vert != root) {
                path.push_back(prev[tmp_vert]);
                tmp_vert = prev[tmp_vert].first;
            }
            std::reverse(path.begin(), path.end());

            ///////////////////////////////
            // add edges in path to graph
            ///////////////////////////////
            // if vertex isn't in graph, add to graph
            if(return_graph.vtoid.find(path[0].first) == return_graph.vtoid.end()) {
                std::pair<uint, bool> add_ret = return_graph.add_unique_vertex(path[0].first);
                assert(add_ret.second);
            }

            for(uint jj=1; jj<path.size(); ++jj) {
                // if vertex isn't in graph, add to graph
                if(return_graph.vtoid.find(path[jj].first) == return_graph.vtoid.end()) {
                    std::pair<uint, bool> add_ret = return_graph.add_unique_vertex(path[jj].first);
                    assert(add_ret.second);
                }
                return_graph.add_edge(path[jj-1].first, path[jj].first, path[jj-1].second, 0, 0);
            }
        }
        // compute number of terminals covered by return_graph
        int num_terminals = 0;
        for(typename Graph<T>::v_to_id::const_iterator viter = return_graph.vtoid.begin(); viter!=return_graph.vtoid.end(); ++viter) {
            if(graph.terminal_set.find(viter->first) != graph.terminal_set.end()) {
                ++num_terminals;
            }
        }
        // return tree
        return return_graph;
    } // end steiner_approx for i == 1
    else { // i > 1
        Graph<T> return_graph;
        while(k > 0) {
            //assert(k == (int)terminal_set.size());
            Graph<T> best_tree;
            double best_cost = std::numeric_limits<double>::infinity();
            for(typename Graph<T>::v_to_id::const_iterator viter = graph.vtoid.begin(); viter != graph.vtoid.end(); ++viter) {
                Graph<T> tmp_tree;
                // compute shortest path from root to viter->first
                std::vector<std::pair<T,char> > path = shortest_path(graph, root, viter->first);
                for(int k2 = 1; k2 <= k; ++k2) {
                    // recursive call
                    if(graph.terminal_set.find(viter->first) != graph.terminal_set.end()) {
                        continue;
                    }
                    tmp_tree = steiner_approx(graph, viter->first, terminal_set, i-1, k2);
                    // compute the number of terminals covered by tmp_tree
                    int num_terminals2 = 0;
                    for(adj_list::const_iterator viter2=tmp_tree.adj.begin(); viter2 != tmp_tree.adj.end(); ++viter2) {
                        // (compute number of terminal vertices covered;
                        // assert that it is equal to k2)
                        if(graph.terminal_set.find(tmp_tree.idtov.at(viter2->first)) != tmp_tree.terminal_set.end()) {
                            ++num_terminals2;
                        }
                    } // end iterating over edges in tree
                    assert(num_terminals2 == k2);

                    // add in path from root to viter->first
                    // if vertex isn't in graph, add to graph
                    if(tmp_tree.vtoid.find(path[0].first) == tmp_tree.vtoid.end()) {
                        std::pair<uint, bool> add_ret = tmp_tree.add_unique_vertex(path[0].first);
                        assert(add_ret.second);
                    }
                    for(uint jj=1; jj<path.size(); ++jj) {
                        // if vertex isn't in graph, add to graph
                        if(tmp_tree.vtoid.find(path[jj].first) == tmp_tree.vtoid.end()) {
                            std::pair<uint, bool> add_ret = tmp_tree.add_unique_vertex(path[jj].first);
                            assert(add_ret.second);
                        }
                        tmp_tree.add_edge(path[jj-1].first, path[jj].first, path[jj-1].second, 0, 0);
                    }

                    ///////////////////////
                    // edges are defind by source, target, and type
                    // need to modify dijkstra's
                    // need to modify shortest paths
                    // need to modify tree construction

                    // compute cost of tree
                    double cost_of_tree = 0.0;
                    int num_terminals = 0;
                    // iterate over edges in tree and add up costs
                    for(adj_list::const_iterator viter=tmp_tree.adj.begin(); viter != tmp_tree.adj.end(); ++viter) {
                        for(edge_list::const_iterator eiter = viter->second.first.begin(); eiter != viter->second.first.end(); ++eiter) {
                            assert(graph.edge_data.size() == graph.num_edges());
                            boost::tuple<char, int, int> edge_d = graph.edge_data.at(boost::make_tuple(tmp_tree.idtov.at(viter->first), tmp_tree.idtov.at(eiter->get<0>()), eiter->get<1>()));
                            cost_of_tree += graph.edge_costs.at(edge_d.get<0>()) * edge_d.get<2>();
                        }
                        // (compute number of terminal vertices covered;
                        // assert that it is equal to k2)
                        if(graph.terminal_set.find(tmp_tree.idtov.at(viter->first)) != tmp_tree.terminal_set.end()) {
                            ++num_terminals;
                        }
                    } // end iterating over edges in tree
                    if(num_terminals != 0) {
                        cost_of_tree = cost_of_tree / num_terminals;
                    } else {
                        cost_of_tree = std::numeric_limits<double>::infinity();
                    }

                    if(cost_of_tree < best_cost) {
                        best_tree = tmp_tree;
                        best_cost = cost_of_tree;
                    }
                } // end iterating over k2
            } // end iterating over vertices

            ///////////////////////////////////////////////////////////
            // add vertices and edges from best tree into return tree
            // iterate over vertices and edges in best_tree
            merge_trees(return_graph, best_tree);

            // compute set of terminals covered by best tree
            boost::unordered_set<T> tree_terminals;
            for(typename Graph<T>::v_to_id::const_iterator viter = best_tree.vtoid.begin(); viter != best_tree.vtoid.end(); ++viter) {
                if(graph.terminal_set.find(viter->first)!=graph.terminal_set.end()) {
                    tree_terminals.insert(viter->first);
                }
            }

            // compute set of terminals covered by return_graph 
            boost::unordered_set<T> return_terminals;
            for(typename Graph<T>::v_to_id::const_iterator viter = return_graph.vtoid.begin(); viter != return_graph.vtoid.end(); ++viter) {
                if(graph.terminal_set.find(viter->first)!=graph.terminal_set.end()) {
                    return_terminals.insert(viter->first);
                }
            }

            // subtract out the number of terminals covered by best tree
            // this is not necessarily equal to k2
            k -= tree_terminals.size();
            //k = k_orig - return_terminals.size(); 
            assert(k >= 0);

            // remove terminal vertices covered by best tree from terminal_set
            for(typename boost::unordered_set<T>::const_iterator viter=tree_terminals.begin(); viter != tree_terminals.end(); ++viter) {
                terminal_set.erase(*viter);
            }
        } // end while loop over k
        //assert(terminal_set.size() == 0);
        // compute set of terminals covered by return graph
        boost::unordered_set<T> tree_terminals;
        for(typename Graph<T>::v_to_id::const_iterator viter = return_graph.vtoid.begin(); viter != return_graph.vtoid.end(); ++viter) {
            if(graph.terminal_set.find(viter->first)!=graph.terminal_set.end()) {
                tree_terminals.insert(viter->first);
            }
        }
        assert((int)tree_terminals.size() == k_orig);

        return return_graph;
    } // end steiner_approx for i > 1
}

template<class T> std::vector<std::pair<T,char> > SteinerApprox::shortest_path(Graph<T> const& graph, T source, T dest) {
    std::set<std::pair<int, uint>, pair_first_less<int, uint> > pq;
    pq.insert(std::pair<int, uint>(0.0, graph.vtoid.at(source)));

    boost::unordered_set<uint> visited;
    visited.insert(graph.vtoid.at(source));

    boost::unordered_map<T, std::pair<T, char> > prev;
    boost::unordered_map<T, double> distances;
    distances[source] = 0.0;

    while(!pq.empty()) {
        uint cur_vert = pq.begin()->second;
        pq.erase(pq.begin());
        visited.insert(cur_vert);

        //if(graph.idtov.at(cur_vert) == dest) {
        //    break;
        //}

        // iterate over outgoing edges
        for(edge_list::const_iterator eiter=graph.adj.at(cur_vert).first.begin(); eiter!=graph.adj.at(cur_vert).first.end(); ++eiter) {
            if(eiter->get<1>() == 4) {
                assert(graph.idtov.at(cur_vert).first == graph.idtov.at(eiter->get<0>()).first && graph.idtov.at(cur_vert).first == graph.idtov.at(eiter->get<0>()).second);
            }
            if(eiter->get<1>() == 5) {
                assert(graph.idtov.at(cur_vert).second == graph.idtov.at(eiter->get<0>()).first && graph.idtov.at(cur_vert).second == graph.idtov.at(eiter->get<0>()).second);
            }
            if(visited.find(eiter->get<0>()) == visited.end()) {
                // compute cur_distance
                // distance to cur_vert + weight*cost
                assert(distances.find(graph.idtov.at(cur_vert)) != distances.end());
                double cur_distance = distances[graph.idtov.at(cur_vert)] + eiter->get<3>() * graph.edge_costs.at(eiter->get<1>());

                if(distances.find(graph.idtov.at(eiter->get<0>())) != distances.end()) {
                    assert(prev.find(graph.idtov.at(eiter->get<0>())) != prev.end());
                    if(distances[graph.idtov.at(eiter->get<0>())] > cur_distance) {
                        // update key for target vertex
                        pq.erase(std::pair<double, uint> (distances[graph.idtov.at(eiter->get<0>())], eiter->get<0>()));
                        distances[graph.idtov.at(eiter->get<0>())] = cur_distance;
                        pq.insert(std::pair<double, uint>(cur_distance, eiter->get<0>()));
                        prev[graph.idtov.at(eiter->get<0>())] = std::pair<T,char>(graph.idtov.at(cur_vert), eiter->get<1>());
                    }
                } else { // target vertex hasn't been explored yet
                    distances[graph.idtov.at(eiter->get<0>())] = cur_distance;
                    pq.insert(std::pair<double, uint>(cur_distance, eiter->get<0>()));
                    prev[graph.idtov.at(eiter->get<0>())] = std::pair<T,char>(graph.idtov.at(cur_vert), eiter->get<1>());
                }
            }
        } // done iterating over outgoing edges from cur_vert
    } // done processing priority queue
    // extract path from source to dest
    std::vector<std::pair<T, char> > tmp_vector;

    tmp_vector.push_back(std::pair<T,char>(dest, -1));
    T tmp_v = dest; 
    
    while(tmp_v != source) {
        tmp_vector.push_back(prev[tmp_v]);
        tmp_v = prev[tmp_v].first;
    }
    std::reverse(tmp_vector.begin(), tmp_vector.end());
    return tmp_vector;
}

template<class T> int SteinerApprox::dijkstra(Graph<T> const& graph, T source, boost::unordered_map<T, double>& distances, boost::unordered_map<T,std::pair<T,char> >& prev) {
    std::set<std::pair<double, uint>, pair_first_less<double, uint> > pq;
    pq.insert(std::pair<double, uint>(0.0, graph.vtoid.at(source)));
    boost::unordered_set<uint> visited;
    visited.insert(graph.vtoid.at(source));

    distances[source] = 0.0;

    while(!pq.empty()) {
        uint cur_vert = pq.begin()->second;
        pq.erase(pq.begin());
        visited.insert(cur_vert);

        // iterate over outgoing edges
        for(edge_list::const_iterator eiter=graph.adj.at(cur_vert).first.begin(); eiter!=graph.adj.at(cur_vert).first.end(); ++eiter) {
            if(visited.find(eiter->get<0>()) == visited.end()) {
                if(eiter->get<1>() == 4) {
                    assert(graph.idtov.at(cur_vert).first == graph.idtov.at(eiter->get<0>()).first && graph.idtov.at(cur_vert).first == graph.idtov.at(eiter->get<0>()).second);
                }
                if(eiter->get<1>() == 5) {
                    assert(graph.idtov.at(cur_vert).second == graph.idtov.at(eiter->get<0>()).first && graph.idtov.at(cur_vert).second == graph.idtov.at(eiter->get<0>()).second);
                }
                // compute cur_distance
                // distance to cur_vert + weight*cost
                assert(distances.find(graph.idtov.at(cur_vert)) != distances.end());
                double cur_distance = distances[graph.idtov.at(cur_vert)] + eiter->get<3>() * graph.edge_costs.at(eiter->get<1>());
                if(distances.find(graph.idtov.at(eiter->get<0>())) != distances.end()) {
                    assert(prev.find(graph.idtov.at(eiter->get<0>())) != prev.end());
                    if(distances[graph.idtov.at(eiter->get<0>())] > cur_distance) {
                        // update key for target vertex
                        pq.erase(std::pair<double, uint> (distances[graph.idtov.at(eiter->get<0>())], eiter->get<0>()));
                        distances[graph.idtov.at(eiter->get<0>())] = cur_distance;
                        pq.insert(std::pair<double, uint>(cur_distance, eiter->get<0>()));
                        prev[graph.idtov.at(eiter->get<0>())] = std::pair<T, char>(graph.idtov.at(cur_vert), eiter->get<1>());
                    }
                } else { // target vertex hasn't been explored yet
                    distances[graph.idtov.at(eiter->get<0>())] = cur_distance;
                    pq.insert(std::pair<double, uint>(cur_distance, eiter->get<0>()));
                    prev[graph.idtov.at(eiter->get<0>())] = std::pair<T,char>(graph.idtov.at(cur_vert), eiter->get<1>());
                }
            } // done iterating over outgoing edges from cur_vert
        }
    } // done processing priority queue
    return 0;
}

template<class T> int SteinerApprox::merge_trees(Graph<T>& graph1, Graph<T> const& graph2) {
    // iterate over vertices in graph2
    for(adj_list::const_iterator viter=graph2.adj.begin(); viter!=graph2.adj.end(); ++viter) {
        // if vertex is not in graph1, then add it
        if(graph1.vtoid.find(graph2.idtov.at(viter->first)) == graph1.vtoid.end()) {
            std::pair<uint, bool> add_ret = graph1.add_unique_vertex(graph2.idtov.at(viter->first));
            assert(add_ret.second);
        }
        // iterate over outgoing edges in graph2
        for(edge_list::const_iterator eiter=viter->second.first.begin(); eiter!=viter->second.first.end(); ++eiter) {
            // if target vertex is not in graph1, then add it
            if(graph1.vtoid.find(graph2.idtov.at(eiter->get<0>())) == graph1.vtoid.end()) {
                std::pair<uint, bool> add_ret = graph1.add_unique_vertex(graph2.idtov.at(eiter->get<0>()));
                assert(add_ret.second);
            }
            // then add edge
            graph1.add_edge(graph2.idtov.at(viter->first), graph2.idtov.at(eiter->get<0>()), eiter->get<1>(), 0, 0);
        }
    }
    return 0;
}

#endif
